﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace Kokomo.Mirror.PeFormat
{
	public class ExportDirectory
	{
		public static ExportDirectory LoadFromModule(
			PeModule module
			)
		{
			if (module == null) throw new ArgumentNullException("module");

			DataDirectoryEntry entry = module.ExportDirectoryEntry;
			if (entry == null)
				throw new ArgumentException("The module does not contain an export directory", "module");

			ExportDirectory directory = new ExportDirectory();
			StructSerializer serializer = new StructSerializer();
			using (Stream stream = module.GetStream())
			{
				stream.Position = module.RvaToFileOffset(entry.VirtualAddress).Value;

				ExportDirectoryHeader exportHeader = serializer.ReadStruct<ExportDirectoryHeader>(stream);

				stream.Position = module.RvaToFileOffset(exportHeader.AddressOfFunctions).Value;
				uint[] functionRvas = serializer.ReadArray<uint>((int)exportHeader.NumberOfFunctions, stream);

				stream.Position = module.RvaToFileOffset(exportHeader.AddressOfNames).Value;
				uint[] nameOffsets = serializer.ReadArray<uint>((int)exportHeader.NumberOfFunctions, stream);

				stream.Position = module.RvaToFileOffset(exportHeader.AddressOfNameOrdinals).Value;
				short[] nameOrdinals = serializer.ReadArray<short>((int)exportHeader.NumberOfNames, stream);

				BinaryReader reader = new BinaryReader(stream);

				Dictionary<int, string> nameOrdinalTable = new Dictionary<int, string>();
				for (int i = 0; i < exportHeader.NumberOfNames; i++)
				{
					short ordinal = nameOrdinals[i];
					uint nameOffset = nameOffsets[i];

					stream.Position = module.RvaToFileOffset(nameOffset).Value;
					string name = reader.ReadNullTerminatedString();
					nameOrdinalTable.Add(ordinal, name);
				}

				for (int i = 0; i < exportHeader.NumberOfFunctions; i++)
				{
					uint functionAddress = functionRvas[i];
					if (functionAddress == 0) continue;

					string name;
					nameOrdinalTable.TryGetValue(i, out name);

					string forwardModuleName = null;
					string forwardFunctionName = null;
					if (entry.Contains(functionAddress))
					{
						stream.Position = module.RvaToFileOffset(functionAddress).Value;
						string forwarder = reader.ReadNullTerminatedString();
						int iExclaim = forwarder.IndexOf('.');
						if (iExclaim > 0)
						{
							forwardModuleName = forwarder.Substring(0, iExclaim);
							forwardFunctionName = forwarder.Substring(iExclaim + 1);
						}
						else
						{
							forwardFunctionName = forwarder;
						}
					}

					ExportedFunction function = new ExportedFunction(
						functionAddress,
						i + exportHeader.Base,
						name,
						forwardModuleName,
						forwardFunctionName
						);
					directory.Exports.Add(function);
				}
			}

			return directory;
		}

		public ExportDirectory()
		{
			this.Exports = new List<ExportedFunction>();
		}

		public IList<ExportedFunction> Exports { get; private set; }
	}
}
